
import glob
from hparams import hparams
from IPython.display import Audio, display, HTML
import librosa
import matplotlib.pyplot as plt
import numpy as np
import os


def load_loss(path):
    steps = []
    losses = []
    for line in open(path).readlines():
        step, loss = line[:-1].split(',')
        step, loss = int(step), float(loss)
        steps.append(step)
        losses.append(loss)
    return np.asarray(steps), np.asarray(losses)

def show_losses(steps, train_losses, test_losses):
    fig, ax = plt.subplots(dpi=100)
    fig.suptitle('Loss curves')
    ax.plot(steps, train_losses, label='Train')
    ax.plot(steps, test_losses, label='Test')
    ax.set_xlabel('Batch')
    ax.set_ylabel('Loss')
    ax.legend()
    plt.show()

def config_mel_ax(ax):
    def format_x(value, tick_number):
        return '%.1f' % (value * hparams.hop_size / hparams.sample_rate)
    
    mel_freqs = librosa.mel_frequencies(n_mels=hparams.num_mels, fmin=hparams.fmin, fmax=hparams.fmax)
    def format_y(value, tick_number):
        i = round(value * (hparams.num_mels - 1))
        return '%.1f' % mel_freqs[i]

    ax.xaxis.set_major_formatter(plt.FuncFormatter(format_x))
    ax.set_xlabel('Time (seconds)')
    
    ax.yaxis.set_major_formatter(plt.FuncFormatter(format_y))
    ax.set_ylabel('Frequency (Hz)')

def show_mel(mel):
    fig, ax = plt.subplots()
    ax.imshow(mel,
              aspect='auto',
              cmap='coolwarm',
              extent=(0, mel.shape[1], 0, 1),
              interpolation='nearest',
              origin='lower')
    config_mel_ax(ax)
    plt.show()

def show_generated_examples(paths):
    gen_paths = []
    for path in paths:
        gen_paths.extend(glob.glob(os.path.join(path, '*_gen.wav')))
    for i, gen_path in enumerate(gen_paths):        
        if i != 0:
            display(HTML('<hr style="border:1px solid black;">'))
        ref_path = gen_path.replace('_gen', '_ref')
        display(HTML('<p>Generated example %d</p>' % (i + 1)))
        display(Audio(filename=gen_path))
        display(HTML('<p>Test set example %d</p>' % (i + 1)))
        display(Audio(filename=ref_path))
        mel_filename = os.path.split(gen_path)[-1].replace('_gen.wav', '-feats.npy')
        mel_paths = glob.glob(os.path.join('org', '*', mel_filename))
        if len(mel_paths) == 0:
            raise ValueError('Mel spectrogram not found')
        elif len(mel_paths) > 1:
            raise ValueError('Multiple mel spectrograms found')
        mel_path = mel_paths[0]
        show_mel(np.load(mel_path).T)


steps, train_losses = load_loss('train_no_dev.txt')
_, test_losses = load_loss('dev.txt')

show_losses(steps, train_losses, test_losses)


show_generated_examples(['generated/checkpoint_latest/dev', 'generated/checkpoint_latest/eval'])